import os 
import json
from multiprocessing import Pool


files = os.listdir("./01.local_density_json")
files_path = [os.path.join("./01.local_density_json", x) for x in files]

output_dir = "./15.density/"
os.makedirs(output_dir, exist_ok=True)

tasks = list(zip(files, files_path, [output_dir] * len(files)))


def generte_dnesity(name, name_path, output_dir):
    with open(name_path, encoding="utf-8") as file:
        content = json.load(file)
    species_all = []
    chromosome = {}
    for k, v in content.items():
        chromosome[k] = v[1]
        species_all.extend(v[1])
    chromosome["all"] = species_all
    with open(output_dir + name.split(".")[0] + "___density.json", mode="w", encoding="utf-8") as file:
        json.dump(chromosome, file, indent=4)
        
        
with Pool(250) as pool:
    pool.starmap(generte_dnesity, tasks)